package com.jorado.elasticsearch;

import com.jorado.core.utility.StringUtils;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.get.MultiGetItemResponse;
import org.elasticsearch.action.get.MultiGetResponse;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.update.UpdateRequestBuilder;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.script.Script;
import org.javatuples.KeyValue;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;

public class ElasticSearchClient {
    private TransportClient client;
    private boolean needClose = false;

    public ElasticSearchClient(TransportClient client) {
        this.client = client;
    }

    public ElasticSearchClient(boolean needClose, TransportClient client) {
        this.needClose = needClose;
        this.client = client;
    }

    public void setClient(TransportClient client) {
        this.client = client;
    }

    public TransportClient getClient() {
        return client;
    }

    /**
     * 连接为空或已经被关闭
     *
     * @return
     */
    public boolean clientIsNullOrClosed() {
        return client == null || client.connectedNodes() == null || client.connectedNodes().size() == 0;
    }

    /**
     * 添加索引
     *
     * @param esIndex
     * @param esType
     * @param doc     doc中 key为: @id   doc主键值 为空时自动创建
     * @return
     */
    public IndexResponse index(String esIndex, String esType, Map<String, Object> doc) {
        if (!clientIsNullOrClosed()) {
            try {
                List<IndexRequestBuilder> indexRequestBuilders = buildIndexRequest(esIndex, esType, doc);
                if (indexRequestBuilders != null && indexRequestBuilders.size() == 1)
                    return indexRequestBuilders.get(0).get();
            } finally {
                if (needClose)
                    client.close();
            }
        }
        return null;
    }

    /**
     * 添加索引
     *
     * @param esIndex
     * @param esType
     * @param docs    doc中 key为: @id   doc主键值 为空时自动创建
     * @return
     */
    public BulkResponse index(String esIndex, String esType, Map<String, Object>[] docs) {
        if (!clientIsNullOrClosed()) {
            try {
                List<IndexRequestBuilder> indexRequestBuilders = buildIndexRequest(esIndex, esType, docs);
                if (indexRequestBuilders.size() > 0) {
                    BulkRequestBuilder bulkRequestBuilder = client.prepareBulk();
                    for (IndexRequestBuilder indexRequestBuilder : indexRequestBuilders)
                        bulkRequestBuilder.add(indexRequestBuilder.request());
                    return bulkRequestBuilder.execute().actionGet();
                }
            } finally {
                if (needClose)
                    client.close();
            }
        }
        return null;
    }

    /**
     * 构建索引请求
     *
     * @param esIndex
     * @param esType
     * @param docs    doc中 key为: _id   doc主键值 为空时自动创建
     * @return
     */
    private List<IndexRequestBuilder> buildIndexRequest(String esIndex, String esType, Map<String, Object>... docs) {
        if (!clientIsNullOrClosed()) {
            try {
                if (docs.length > 0) {
                    List<IndexRequestBuilder> indexRequests = new ArrayList<>();
                    for (Map<String, Object> doc : docs) {
                        IndexRequestBuilder indexRequestBuilder;
                        Object id = doc.containsKey("_id") ? doc.remove("_id") : null;
                        if (id == null || StringUtils.isNullOrEmpty(id.toString()))
                            indexRequestBuilder = client.prepareIndex(esIndex, esType);
                        else
                            indexRequestBuilder = client.prepareIndex(esIndex, esType, id.toString());
                        XContentBuilder xContentBuilder = jsonBuilder().startObject();
                        for (String key : doc.keySet()) {
                            xContentBuilder.field(key, doc.get(key));
                        }
                        xContentBuilder.endObject();
                        indexRequestBuilder.setSource(xContentBuilder).request();
                        indexRequests.add(indexRequestBuilder);
                    }
                    return indexRequests;
                }
            } catch (Exception ex) {

            }
        }
        return null;
    }

    /**
     * 更新索引
     *
     * @param esIndex
     * @param esType
     * @param upsert
     * @param doc     doc中 key为: _id   doc主键值
     * @return
     */
    public UpdateResponse update(String esIndex, String esType, Boolean upsert, Map<String, Object> doc) {
        if (!clientIsNullOrClosed()) {
            try {
                List<UpdateRequestBuilder> updateRequestBuilders = buildUpdateRequest(esIndex, esType, upsert == null ? false : upsert, doc);
                if (updateRequestBuilders != null && updateRequestBuilders.size() == 1) {
                    return updateRequestBuilders.get(0).get();
                }
            } finally {
                if (needClose)
                    client.close();
            }
        }
        return null;
    }

    /**
     * 更新索引
     *
     * @param esIndex
     * @param esType
     * @param upsert
     * @param docs    doc中 key为: @id   doc主键值
     * @return
     */
    public BulkResponse update(String esIndex, String esType, Boolean upsert, Map<String, Object>... docs) {
        if (!clientIsNullOrClosed()) {
            try {
                List<UpdateRequestBuilder> updateRequestBuilders = buildUpdateRequest(esIndex, esType, upsert == null ? false : upsert, docs);
                if (updateRequestBuilders != null && updateRequestBuilders.size() > 0) {
                    BulkRequestBuilder bulkRequestBuilder = client.prepareBulk();
                    for (UpdateRequestBuilder updateRequestBuilder : updateRequestBuilders)
                        bulkRequestBuilder.add(updateRequestBuilder);
                    return bulkRequestBuilder.execute().actionGet();
                }
            } finally {
                if (needClose)
                    client.close();
            }
        }
        return null;
    }

    /**
     * @param esIndex
     * @param esType
     * @param upsert
     * @param idScrpit
     * @return
     */
    public UpdateResponse update(String esIndex, String esType, Boolean upsert, KeyValue<String, Script> idScrpit) {
        if (!clientIsNullOrClosed()) {
            try {
                List<UpdateRequestBuilder> updateRequestBuilders = buildUpdateRequest(esIndex, esType, upsert, idScrpit);
                if (updateRequestBuilders != null && updateRequestBuilders.size() == 1) {
                    return updateRequestBuilders.get(0).get();
                }
            } finally {
                if (needClose)
                    client.close();
            }
        }
        return null;
    }

    /**
     * @param esIndex
     * @param esType
     * @param upsert
     * @param idScripts
     * @return
     */
    public BulkResponse update(String esIndex, String esType, Boolean upsert, KeyValue<String, Script>... idScripts) {
        if (!clientIsNullOrClosed()) {
            try {
                List<UpdateRequestBuilder> updateRequestBuilders = buildUpdateRequest(esIndex, esType, upsert, idScripts);
                if (updateRequestBuilders != null && updateRequestBuilders.size() > 0) {
                    BulkRequestBuilder bulkRequestBuilder = client.prepareBulk();
                    for (UpdateRequestBuilder updateRequestBuilder : updateRequestBuilders)
                        bulkRequestBuilder.add(updateRequestBuilder);
                    return bulkRequestBuilder.execute().actionGet();
                }
            } finally {
                if (needClose)
                    client.close();
            }
        }
        return null;
    }

    /**
     * 构建更新请求
     *
     * @param esIndex
     * @param esType
     * @param upsert
     * @param docs    doc中 key为: _id   doc主键值
     * @return
     */
    private List<UpdateRequestBuilder> buildUpdateRequest(String esIndex, String esType, Boolean upsert, Map<String, Object>... docs) {
        if (!clientIsNullOrClosed() && docs.length > 0) {
            List<UpdateRequestBuilder> updateRequestBuilders = new ArrayList<>();
            for (Map<String, Object> doc : docs) {
                Object id = doc.containsKey("_id") ? doc.remove("_id") : null;
                if (id == null || StringUtils.isNullOrEmpty(id.toString()))
                    continue;
                UpdateRequestBuilder requestBuilder = client.prepareUpdate(esIndex, esType, id.toString()).setDoc(doc);
                if (upsert)
                    requestBuilder.setDocAsUpsert(upsert);
                updateRequestBuilders.add(requestBuilder);
            }
            return updateRequestBuilders;
        }
        return null;
    }

    /**
     * @param esIndex
     * @param esType
     * @param upsert
     * @param idScripts KEY:id,VALUE:script
     * @return
     */
    private List<UpdateRequestBuilder> buildUpdateRequest(String esIndex, String esType, Boolean upsert, KeyValue<String, Script>... idScripts) {
        if (!clientIsNullOrClosed() && idScripts != null) {
            List<UpdateRequestBuilder> updateRequestBuilders = new ArrayList<>();

            for (KeyValue<String, Script> keyValue : idScripts) {
                if (keyValue == null || StringUtils.isNullOrEmpty(keyValue.getKey()) || keyValue.getValue() == null)
                    continue;
                UpdateRequestBuilder requestBuilder = client.prepareUpdate(esIndex, esType, keyValue.getKey()).setScript(keyValue.getValue());
                if (upsert) {
                    requestBuilder.setScriptedUpsert(true);
                    requestBuilder.setUpsert(new HashMap<>());
                }
                updateRequestBuilders.add(requestBuilder);
            }
            return updateRequestBuilders;
        }
        return null;
    }

    /**
     * 根据ID获取文档信息
     *
     * @param esIndex
     * @param esType
     * @param ids
     * @return
     */
    public List<String> getDocumentStringById(String esIndex, String esType, String... ids) {
        List<String> result = new ArrayList<>();
        if (!clientIsNullOrClosed()) {
            try {
                MultiGetResponse responses = client.prepareMultiGet().add(esIndex, esType, ids).get();
                if (responses != null && responses.getResponses() != null)
                    for (MultiGetItemResponse response : responses.getResponses()) {
                        result.add(response.getResponse().getSourceAsString());
                    }
            } finally {
                if (needClose)
                    client.close();
            }
        }
        return result;
    }

    /**
     * 根据ID获取文档信息
     *
     * @param esIndex
     * @param esType
     * @param ids
     * @return
     */
    public List<Map<String, Object>> getDocumentMapById(String esIndex, String esType, String... ids) {
        List<Map<String, Object>> result = new ArrayList<>();
        if (!clientIsNullOrClosed()) {
            try {
                MultiGetResponse responses = client.prepareMultiGet().add(esIndex, esType, ids).get();
                if (responses != null && responses.getResponses() != null)
                    for (MultiGetItemResponse response : responses.getResponses()) {
                        result.add(response.getResponse().getSourceAsMap());
                    }
            } finally {
                if (needClose)
                    client.close();
            }
        }
        return result;
    }
}
